import gymnasium as gym
from CliffWalking import PolicyIteration, rollout_gridworld

# 创建 FrozenLake v1 环境（用 ansi 模式在控制台输出地图）
env = gym.make("FrozenLake-v1", render_mode="ansi", is_slippery=False)
env = env.unwrapped  # 解封装才能访问 P

state, info = env.reset()   # 必须 reset 一下，初始化状态
print("""
S：起点 (Start, 索引 0)。
F：安全的冰面 (Frozen)。
H：冰洞 (Hole)，掉进去游戏结束。
G：目标 (Goal)，到达后奖励 1。
""")
print(env.render())         # 渲染地图，返回字符串，打印出来

holes = set()
ends = set()
for s in env.P:
    for a in env.P[s]:
        for s_ in env.P[s][a]:
            if s_[2] == 1.0:  # 获得奖励为1,代表是目标
                ends.add(s_[1])
            if s_[3] is True:
                holes.add(s_[1])
holes = holes - ends

print("冰洞的索引:", holes)
print("目标的索引:", ends)

# 查看遍历所有状态的状态转移信息
for s in env.P:       # 遍历所有状态
    for a in env.P[s]:  # 遍历该状态的所有动作
        print(s, a, env.P[s][a])


# 这个动作意义是Gym库针对冰湖环境事先规定好的
action_meaning = ['<', 'v', '>', '^']
theta = 1e-5
gamma = 0.9
agent = PolicyIteration(env, theta, gamma)
agent.policy_iteration(action_meaning, [5, 7, 11, 12], [15])

holes = {5, 7, 11, 12}
goal = 15
rollout_gridworld(agent.pi, nrow=4, ncol=4, start_state=0, goal_state=goal, fail_states=holes, n_episodes=1000)




# action_meaning = ['<', 'v', '>', '^']
# theta = 1e-5
# gamma = 0.9
# agent = ValueIteration(env, theta, gamma)
# agent.value_iteration()
# print_agent(agent, action_meaning, [5, 7, 11, 12], [15])